# -*- coding: utf-8 -*-
"""
@author:xuyuntao
@time:2021/5/4:11:34
@email:xuyuntao@189.cn
"""
import numpy
import matplotlib.font_manager
import matplotlib.pyplot as plt
import seaborn as sns
from Tools import getLevels
from tqdm import tqdm

"""用于测试各种输入信号既可以出一个标量值的特征量，修改时
注意添加specialList列表中列出的需要特数输入参数的函数"""
# -------第一幅图绘制热力图-------
def drawColorMap(figNum,data,xLabel="",yLabel="",xtickLabels=None,ytickLabels=None,
                 title=None,level_min=None,level_max=None,
                 colormap="nipy_spectral",figName=None,
                 figSizeNum = 8,figScale=(4,3),
                 lineWidth=1,numbersInsideSize=8,
                 labelSize = 18,fontSize = 21):
    """绘制colormap，各参数命名与函数中命名一致"""
    zhFont = matplotlib.font_manager.FontProperties(fname='C:\Windows\Fonts\simsun.ttc')  # 中文字体
    figSize = (figSizeNum, round(figSizeNum / figScale[0] * figScale[1]))
    if (level_min==None or level_max==None):
        minLevel, maxLevel = getLevels(data.min(), data.max(), 10)  # 计算参考轴最大最小值
    else:
        minLevel=float(level_min)
        maxLevel=float(level_max)
    fig = plt.figure(figNum, figsize=figSize)  # 生成fig
    ax = fig.add_subplot(1, 1, 1)  # 生成axes
    if xtickLabels==None:
        xtickLabels=[""]*data.shape[1]
    if ytickLabels==None:
        ytickLabels=[""]*data.shape[0]
    ax=sns.heatmap(data, ax=ax,vmin=minLevel,vmax=maxLevel,cmap=colormap,annot=True,
                   linewidths=lineWidth,cbar=False,
                   annot_kws={"fontsize":numbersInsideSize},
                   xticklabels=xtickLabels,
                   yticklabels=ytickLabels)  # 热力图绘制
    ax.tick_params(axis="both", labelsize=labelSize)  # 修改热力图坐标轴编号
    colorbar = ax.figure.colorbar(ax.collections[0])  # 显示colorbar
    colorbar.ax.tick_params(labelsize=labelSize)  # 设置colorbar刻度字体大小。
    if figName==None:
        pass
    else:
        fig.canvas.manager.set_window_title(figName)
    if xLabel==None:
        pass
    else:
        ax.set_xlabel(xLabel, fontproperties=zhFont, fontsize=fontSize)
    if yLabel==None:
        pass
    else:
        ax.set_ylabel(yLabel, fontproperties=zhFont, fontsize=fontSize)

    if title==None:
        pass
    else:
        fig.suptitle(title,fontproperties=zhFont, fontsize=fontSize)  # 标题
    return fig

# -------第二幅图绘制折线图-------
def drawLine(figNum,xData,yData,xLabel="",yLabel="",
             legend=None,xlim=None,ylim=None,figName=None,
             marker=True,markersize = 5,
             figSizeNum = 8,figScale=(4,3),
             linestyle="-",lineWidth = 1,
             legendSize = 21,labelSize = 18,fontSize = 21,framealpha=0.7):
    """绘制曲线图，各参数命名与函数命名一致，第二、三个参数为绘制的x轴数据和y轴数据，
    可以输入多行，则绘制多条曲线，曲线legend可通过legend参数传入"""
    zhFont = matplotlib.font_manager.FontProperties(fname='C:\Windows\Fonts\simsun.ttc')  # 中文字体
    allMarkers=['.', ',', 'o', 'v', '^', '<', '>', '1', '2',
                '3', '4', 's', 'p', '*', 'h', 'H', '+', 'x',
                'D', 'd', '|', '_']
    if xData.shape==xData.shape:
        pass
    else:
        raise ValueError("两Data要形状一直")
    useMarker=None
    ifMarker=False
    if type(marker)==bool:
        if marker:
            if yData.shape[0]>len(allMarkers):
                useMarker=allMarkers.copy()
                useMarker+=[""]*(yData.shape[0]-len(allMarkers))
            else:
                useMarker=allMarkers[:yData.shape[0]].copy()
            ifMarker=True
        else:
            ifMarker=False
    else:
        if marker==None:
            ifMarker = False
        else:
            if type(marker)==list:
                useMarker=marker[:yData.shape[0]].copy()
                ifMarker = True
            else:
                ifMarker = False
    figSize = (figSizeNum, round(figSizeNum / figScale[0] * figScale[1]))
    fig = plt.figure(figNum, figsize=figSize)
    ax = fig.add_subplot(1, 1, 1)
    if ifMarker:
        for lineIndex in range(yData.shape[0]):
            ax.plot(xData[lineIndex, :], yData[lineIndex, :],
                    marker=useMarker[lineIndex],markersize=markersize,
                    linewidth=lineWidth,linestyle=linestyle)
    else:
        for lineIndex in range(yData.shape[0]):
            ax.plot(xData[lineIndex,:],yData[lineIndex,:],linewidth=lineWidth)
    if xlim==None:
        pass
    else:
        ax.set_xlim(xlim[0],xlim[1])
    if ylim==None:
        pass
    else:
        ax.set_ylim(ylim[0],ylim[1])
    if legend==None:
        pass
    else:
        if len(legend)==xData.shape[0]:
            ax.legend(legend,fontsize=legendSize,framealpha=framealpha)
        else:
            raise ValueError("图例和绘图线数长度不一致")
    if xLabel==None:
        pass
    else:
        ax.set_xlabel(xLabel, fontproperties=zhFont, fontsize=fontSize)
    if yLabel==None:
        pass
    else:
        ax.set_ylabel(yLabel, fontproperties=zhFont, fontsize=fontSize)
    ax.tick_params(axis="both", labelsize=labelSize)
    if figName==None:
        pass
    else:
        fig.canvas.manager.set_window_title(figName)
    return fig
